# -*- coding: utf-8 -*-
# SSHexec 异步执行文件

# 系统或第三方模块
import asyncio
from time import sleep
import asyncssh
import threading
import queue
import signal
import re

from datetime import datetime
from typing import Dict, List, Optional
from posixpath import join as posix_join

# 自定义模块
import src.config as config
import src.output as output
import src.utils as utils
import src.asyncssh.asyncssh_utils as asyncssh_utils
from src.utils import tlog, elog, JumpOut 


# 全局结果队列，用于生产者-消费者模式
result_queue = queue.Queue()


async def asyncssh_connect_execute(
    runner: Dict[str, str], 
    command: str, 
    full_command: str,
    connect_timeout: int, 
    exec_timeout: int, 
    env: Dict[str, str] = {},
    use_sudo: bool = False,
    ) -> Dict[str, Optional[str]]:
    """
    功能：
        异步SSH连接执行任务，连接成功后执行命令，将结果输出到结果队列
        
    参数：
        runner: 目标节点信息字典，包含ip, port, user, password
        command: 要执行的命令
        full_command: 要执行的命令，包含sudo和转义引号
        connect_timeout: 连接超时时间
        exec_timeout: 执行超时时间
        env: 环境变量字典
        use_sudo: 是否使用sudo执行

    返回：
        Dict[str, Optional[str]]: 包含执行结果的字典
    """
    thread_start_time = datetime.now()              # 线程开始时间
    shell_start_time = None                         # shell开始时间
    elog.debug(f"[asyncssh]{runner.get('ip')}：线程开始时间：{thread_start_time}")

    result = {
        "ip": runner.get("ip"),                     # 目标节点IP
        "connect_bool": None,                       # 连接状态：是否成功
        "connect_str": None,                        # 连接结果内容
        "exit_bool": None,                          # 执行状态：是否成功
        "exit_status": None,                        # 退出状态码
        "result_category": None,                    # 结果分类

        "command": command,                         # 执行命令
        "stdout": None,                             # 标准输出 含 标准错误输出

        "thread_start_time": thread_start_time,     # 线程开始时间 
        "thread_end_time": None,                    # 线程结束时间
        "thread_cost_time": -1,                     # 执行耗时
        "shell_cost_time": -1,                      # shell执行耗时
        
        "port": runner.get("port"),                 # 目标节点端口
        "user": runner.get("user"),                 # 目标节点用户名 
        "password": runner.get("password"),         # 目标节点密码
        
        "full_command": full_command,               # 执行命令，包含sudo和转义引号
        "env": env,                                 # 环境变量
        "use_sudo": use_sudo,                       # 是否使用sudo执行
        "connect_timeout": connect_timeout,         # 连接超时时间
        "exec_timeout": exec_timeout,               # 执行超时时间
        
    }

    #检查是否收到中断信号
    if asyncssh_utils.asyncssh_interrupted: 
        elog.info(f"[asyncssh]{runner.get('ip')}：检测到中断信号，跳出，不执行后续命令（连接前）")
        result.update({
            "connect_bool": False,
            "connect_str": "收到中断信号",
            "stdout": "收到中断信号",
            "exit_status": -130,
            "exit_bool": False,
            "result_category": "收到中断信号",
        })
        return result

    try:
        # 创建异步协程，建立SSH连接 as conn
        async with await asyncssh_utils.connect_with_retry(runner, connect_timeout) as conn:
            result["connect_str"] = "连接成功"
            result["connect_bool"] = True
            elog.success(f"[asyncssh]{runner.get('ip')}：SSH连接成功，主机：{runner['ip']}，端口：{runner['port']}，用户名：{runner['user']}")



            # 检查是否收到中断信号（阻塞连接后）
            if asyncssh_utils.asyncssh_interrupted: raise JumpOut("收到中断信号，跳出，不执行后续命令（连接后）")
            # 防止连接成功，但未进入shell环境的检测处理
            try:
                elog.info(f"[asyncssh]{runner.get('ip')}：执行“空”命令，检查是否进入shell环境")
                # 执行空命令
                connect_result = await conn.run(':', timeout=10, check=False, stdout=asyncssh.PIPE, stderr=asyncssh.STDOUT)
                elog.info(f"[asyncssh]{runner.get('ip')}：执行“空”命令成功")

                # 处理重定向的标准输出
                if connect_result.stdout:
                    # 处理标准输出
                    command_stdout = connect_result.stdout
                    if isinstance(command_stdout, bytes):
                        command_stdout = command_stdout.decode('utf-8').strip()
                    else:
                        command_stdout = command_stdout.strip()

                    # 密码过期列表
                    passwd_error = ["password expired", "change your password now", "has expired", "密码过期"]
                    # 检查密码过期关键词
                    if any(error in command_stdout.lower() for error in passwd_error):
                        elog.warning(f"[asyncssh]{runner.get('ip')}：“空”命令检测到密码过期关键词")
                        result["result_category"] = "密码已过期"
                        raise Exception("密码已过期")

                elog.info(f"[asyncssh]{runner.get('ip')}：执行“空”命令，退出状态码：{connect_result.exit_status}，标准输出：{connect_result.stdout.strip() if connect_result.stdout else '无标准输出'}")

            except Exception as e:
                elog.error(f"[asyncssh]{runner.get('ip')}：执行空命令异常\n异常类型: {type(e)} \n异常参数: {e.args} \n异常信息: {e}")
            


            # 检查是否收到中断信号（阻塞空命令后）
            if asyncssh_utils.asyncssh_interrupted: raise JumpOut("收到中断信号，跳出，不执行后续命令（空命令后）")
            # 执行命令 记录执行开始前后时间
            shell_start_time = datetime.now()
            elog.info(f"[asyncssh]{runner.get('ip')}：”拼接”命令开始执行，时间：{shell_start_time}")
            
            # 开始执行拼接后的命令
            try:
                command_result = await conn.run(full_command, check=False, stdout=asyncssh.PIPE, stderr=asyncssh.STDOUT)
                # 处理标准输出
                command_stdout = command_result.stdout
                if command_stdout is None:
                    command_stdout = ""
                elif isinstance(command_stdout, bytes):
                    command_stdout = command_stdout.decode('utf-8').strip()
                else:
                    command_stdout = command_stdout.strip()

                # 处理移除clean_stdout中的颜色转义内容 例如：^[[0;33;40m3.09^[[0m 等
                clean_stdout = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])').sub('', command_stdout)
                elog.success(f"[asyncssh]{runner.get('ip')}：命令执行结果，退出状态码：{command_result.exit_status}，标准输出：\n{clean_stdout if clean_stdout else '无标准输出'}")

                result["stdout"] = clean_stdout
                result["exit_status"] = command_result.exit_status
                result["exit_bool"] = command_result.exit_status == 0
                
                # 执行结果分类
                # 初始化退出状态码字典
                exit_code_dict = {
                    0: "执行成功",
                    # 1: "执行失败",
                    # 2: "参数不存在",
                    124: "执行超时",
                    # 125: "命令语法错误",
                    # 126: "命令权限不足",
                    # 127: "命令不存在",
                    # 128: "无效退出码",
                    130: "被中断（Ctrl+C）",
                    137: "被强制终止（SIGKILL）",
                }
                # 从字典中获取执行结果分类，若不存在则赋值为返回值
                if result["exit_status"] is not None:
                    result["result_category"] = exit_code_dict.get(result["exit_status"], f'退出码{result["exit_status"]}')
                else:
                    result["result_category"] = "未获取到退出码"

            except Exception as e:
                elog.error(f"[asyncssh]{runner.get('ip')}：命令执行过程中发生异常\n异常类型: {type(e)} \n异常参数: {e.args} \n异常信息: {e}")
                # 协程处理执行中的异常的错误分类
                result["exit_status"] = -999
                result["result_category"] = "代码执行异常"


    except KeyboardInterrupt:
        elog.warning(f"[asyncssh]{runner.get('ip')}：中断信号已被except捕获")
    
    except JumpOut as e:
        elog.warning(e)
        if asyncssh_utils.asyncssh_interrupted:
            elog.warning(f"[asyncssh]{runner.get('ip')}：收到中断信号，进行中断处理(interrupt_execute)")
            result["connect_bool"] = False if not result["connect_bool"] else result["connect_bool"]
            result["connect_str"] = "收到中断信号" if not result["connect_str"] else result["connect_str"]
            result["stdout"] = "收到中断信号" if not result["stdout"] else result["stdout"]
            # 130（128 + 信号编号 = 退出码） 是Unix/Linux中SIGINT(Ctrl+C)的标准退出码， SIGINT 的信号编号是2，所以 128 + 2 = 130 ，负数表示人为赋值，不与正常的退出码冲突
            result["exit_status"] = -130 if not result["exit_status"] else result["exit_status"]
            result["exit_bool"] = False if not result["exit_bool"] else result["exit_bool"]
            result["result_category"] = "收到中断信号" if not result["result_category"] else f'{result["result_category"]}，收到中断信号'


    except Exception as e:
        elog.error(f"[asyncssh]{runner.get('ip')}：连接过程中发生异常\n异常类型: {type(e)} \n异常参数: {e.args} \n异常信息: {e}")
        # 连接超时会抛出异常，这里捕获异常并分类为连接超时
        if "timeouterror" in str(type(e)).lower():
            e = "连接超时"
        # 协程处理连接中的异常的错误分类
        result["connect_bool"] = False if not result["connect_bool"] else result["connect_bool"]
        result["exit_bool"] = False
        result["connect_str"] = str(e) if not result["connect_str"] else result["connect_str"]
        result["result_category"] = utils.error_classify(runner.get('ip', '未知IP'), str(e)) if not result["result_category"] else result["result_category"]

    finally:
        elog.info(f"[asyncssh]{runner.get('ip')}：运行结果分类，结果为：{result['result_category']}")
        
        if shell_start_time:
            shell_end_time = datetime.now()
            result["shell_cost_time"] = round((shell_end_time - shell_start_time).total_seconds(), 3)
            elog.info(f"[asyncssh]{runner.get('ip')}：”拼接”命令执行结束，时间：{shell_end_time}，耗时：{result['shell_cost_time']}秒")
        
        # 记录执行结束时间和耗时
        thread_end_time = datetime.now()
        elog.debug(f"[asyncssh]{runner.get('ip')}：线程执行结束时间：{thread_end_time}")
        result["thread_end_time"] = thread_end_time
        result["thread_cost_time"] = f"{round((thread_end_time - thread_start_time).total_seconds(), 3)}"
        elog.info(f"[asyncssh]{runner.get('ip')}：线程执行耗时（秒）：{result['thread_cost_time']}")

        # 结果放入消费者队列
        result_queue.put(result)
        elog.info(f"[asyncssh]{runner.get('ip')}：线程执行结束，已将结果放入消费者队列")
        return result
    

async def create_asyncio_task(
    runners: List[Dict[str, str]], 
    env: Dict[str, str], 
    use_sudo: bool,
    command: str, 
    full_command: str,
    connect_timeout: int, 
    exec_timeout: int, 
    max_concurrency: int,
    ) -> List[Dict[str, Optional[str]]]:
    """
    功能：
        创建异步任务，用于并发执行远程命令

    参数：
        runners: 目标节点信息列表
        env: 环境变量
        use_sudo: 是否使用sudo
        command: 要执行的命令
        full_command: 完整的命令字符串
        connect_timeout: 连接超时时间
        exec_timeout: 执行超时时间
        max_concurrency: 最大并发数

    返回：
        List[Dict[str, Optional[str]]]: 所有节点的执行结果列表
    """

    # 使用信号量控制任务并发数
    semaphore = asyncio.Semaphore(max_concurrency)

    async def _run_with_semaphore(runner):
        # 尝试获取信号量：成功获取（计数器>0）：立即执行execute_remote_command失败获取（计数器=0）：挂起当前协程，进入等待队列
        async with semaphore:
            # 执行远程命令
            return await asyncssh_connect_execute(
                runner, command, full_command, connect_timeout, exec_timeout, env, use_sudo
            )
    
    # 循环字典列表，创建所有任务，创建len(runners)个task任务（创建一组协程对象）， run_with_semaphore(runner)调用返回一个协程对象
    tasks = [_run_with_semaphore(runner) for runner in runners]
    elog.info(f"asyncio：已创建 {len(tasks)} 个任务")
    
    # 并发执行所有tasks任务，未激活状态的任务会被激活，激活后会进入事件循环，等待信号量释放，释放后会再次进入事件循环，直到所有任务都执行完毕
    results = await asyncio.gather(*tasks, return_exceptions=True)
    elog.info(f"asyncio：已激活 {len(results)} 个任务")
    
    # 结果数据字典列表
    final_results = []

    # 处理结果 放入enumerate(results)是 Python 中用于 ​​同时获取迭代元素的索引和值​​ 的高效工具
    elog.info(f"处理结果，结果数量：{len(results)}")
    for i, result in enumerate(results):
        final_results.append(result)

    return final_results

def output_processor(log_dir: str):
    """
    功能：
        输出处理器 - 消费者函数，从队列中取出结果并输出
    
    参数：
        log_dir: 输出文件路径
    """

    # 配置日志文件路径
    log_file = posix_join(log_dir, config.DEFAULT_FILE_OUTPUT)
    
    # 打开输出文件
    with open(log_file, 'w', encoding='utf-8') as f:
        while True:
            elog.info(f"日志消息队列(output_processor):当前等待处理数量: {result_queue.qsize()}")
            try:
                # 从队列中获取结果
                result = result_queue.get(timeout=3)
                # 如果收到结束信号，且队列为空，退出循环
                if result == "DONE":
                    break
                
                # 格式化输出
                output_text = output.format_terminal_output(result)
                
                # 输出到终端
                print(output_text)
                
                # 输出到文件
                f.write(output_text + "\n" + "═" * 48 + "\n")
                f.flush()  # 确保立即写入文件
                
            except queue.Empty:
                # 队列为空，继续等待
                sleep(0.2)
                elog.debug(f"队列为空，继续等待")
                continue
    


async def asyncssh_executor_manager(args, runners: List[Dict[str, str]], log_dir: str) -> List[Dict[str, Optional[str]]]: 
    """
    功能：
        异步执行ssh命令

    参数：
        args: 命令行参数对象
        runners: 目标节点信息列表
        log_dir: 日志目录

    返回：
        final_results: 所有节点的执行结果列表
    """

    # 检查是否需要sudo
    use_sudo = True if args.m == "sudo" else False

    # 仅用于后面字典赋值使用，不参与其他逻辑
    command = args.c if args.c else args.s

    # 拼接要执行的命令
    full_command = asyncssh_utils.add_env_sudo_to_commands(args)
                    
    # 1. 启动输出线程（非守护线程）
    output_thread = threading.Thread(target=output_processor, args=(log_dir,))
    elog.info(f"终端输出线程初始化(output_processor)，日志文件路径：{log_dir}")

    # output_thread.daemon = True
    output_thread.start()
    elog.success(f"终端输出线程启动成功")
    
    try:
        # 2. 创执行远程命令
        elog.info(f"创建远程执行任务开始，共 {len(runners)} 个节点")
        final_results = await create_asyncio_task(
            runners=runners,
            command=command,
            full_command=full_command,
            env=args.e,
            use_sudo=use_sudo,
            connect_timeout=args.T,
            exec_timeout=args.t,
            max_concurrency=args.n,
        )
        elog.success(f"所有节点执行完成，共 {len(final_results)} 个")

    except Exception as e:
        elog.error(f"执行过程中出现异常(asyncssh_executor_manager)\n异常类型: {type(e)} \n异常参数: {e.args} \n异常信息: {e}")
        print(f"{config.COLOR_RED}[WARNING]{config.COLOR_RESET}{config.COLOR_YELLOW}[function:asyncssh_executor_manager]{config.COLOR_RESET} 警告：执行过程中出现异常！异常信息：{str(e)}")
        raise

    finally:
        # 3. 发送结束信号
        result_queue.put("DONE")
        elog.info(f"发送结束信号到队列，已完成")
        
        # 4. 等待消费者线程完成（超时保护）（最多等待10秒）
        output_thread.join(timeout=10)  
        if output_thread.is_alive():
            elog.warning(f"输出线程未在超时时间内完成，已超时10秒")
            print(f"{config.COLOR_RED}[WARNING]{config.COLOR_RESET}{config.COLOR_YELLOW}[function:asyncssh_executor_manager]{config.COLOR_RESET} 警告：消费者线程未在超时时间内完成，已超时10秒！")

        if asyncssh_utils.asyncssh_interrupted:
            elog.success(f"中断清理已完成")
            print(f"{config.COLOR_RED}中断清理已完成{config.COLOR_RESET}")

        elog.success(f"线程任务全部完成，正常返回results")
        return final_results

def execute_asyncssh(args, nodesinfos, log_dir):
    # 将asyncio模块加载到确认执行后，优化main文件启动速度

    # 注册全局中断信号处理器
    signal.signal(signal.SIGINT, asyncssh_utils.handle_interrupt)
    signal.signal(signal.SIGTERM, asyncssh_utils.handle_interrupt)
    elog.success(f"全局中断信号处理器注册成功，信号：SIGINT, SIGTERM")
    
    elog.info(f"开始执行asyncssh，参数：{args}")
    final_results = asyncio.run(asyncssh_executor_manager(args, nodesinfos, log_dir))
    return final_results
